[TRTLLM-11958][perf] reduce @torch.library.custom_op host overhead - #13149

@coderabbitai summary

Description

@torch.library.custom_op wraps every call in Python-level wrappers registered on several dispatch keys (Autograd, optionally ADInplaceOrView, and the backend key), which imposes a ~7us per-call dispatcher tax on top of the actual kernel launch. On host-heavy iterations (LTX-2 dense transformer issues ~2100 such calls per step for the two ops in this PR) the tax becomes a measurable fraction of the per-step wall time.

This PR:

1. Introduces a thin helper fast_custom_op that registers an op directly through the low-level torch.library.Library.define + impl API. The helper preserves the @custom_op developer experience — schema is inferred from Python type hints via torch.library.infer_schema, and the returned object exposes .register_fake — while bypassing the multi-layer Python wrappers that @custom_op installs.
2. Migrates two hot ops on the LTX-2 VisualGen path (trtllm::nvfp4_gemm and trtllm::tunable_fp4_quantize) to @fast_custom_op.

Approach — usage

# Before
@torch.library.custom_op("trtllm::nvfp4_gemm", mutates_args=())
def nvfp4_gemm(act_fp4: torch.Tensor, ...) -> torch.Tensor: ...
@nvfp4_gemm.register_fake
def _(...): ...

# After (same ergonomics, ~5us/call cheaper)
from tensorrt_llm._torch.custom_ops.fast_custom_op import fast_custom_op

@fast_custom_op("trtllm::nvfp4_gemm", mutates_args=())
def nvfp4_gemm(act_fp4: torch.Tensor, ...) -> torch.Tensor: ...
@nvfp4_gemm.register_fake
def _(...): ...

The helper uses torch.library.infer_schema internally so the schema is still driven by Python type hints — no hand-written schema strings. FRAGMENT mode is used under the hood because the trtllm namespace is already declared by C++ via TORCH_LIBRARY_FRAGMENT.

Why @custom_op is expensive — code-level analysis

Looking at torch/_library/custom_ops.py::CustomOpDef._register_to_dispatcher (L607-675), @custom_op registers several Python-level kernels on multiple dispatch keys. Each one is executed on every call:

1. Autograd key (always registered) — torch/_library/autograd.py::autograd_impl (L108):

def autograd_impl(keyset, *args, **keyword_only_args):
    if _C.is_grad_enabled() and _C._any_requires_grad(*args):   # tensor iteration
        result = Generated.apply(*args, Metadata(keyset, kwargs))
    else:
        result = forward_no_grad(*args, Metadata(keyset, kwargs))
    return result

def forward_no_grad(*args):
    metadata = args[-1]; args = args[:-1]
    with _C._AutoDispatchBelowAutograd():                       # Python ctx mgr
        return op.redispatch(                                    # re-dispatch
            keyset & _C._after_autograd_keyset, *args, **metadata.keyword_only_args)

Every call pays: is_grad_enabled() + _any_requires_grad(*args) tensor iteration + Metadata dataclass construction + _AutoDispatchBelowAutograd context manager + op.redispatch(...) (a second dispatch trip).

2. ADInplaceOrView key — adinplaceorview_impl (L654). Registered only when the schema is mutable or a view op. Bumps version counters on mutated args and routes through call_boxed. Not a cost for the two ops in this PR (both are pure, mutates_args=()), so schema is non-mutable and this wrapper is not installed.

3. CUDA backend key — backend_impl wrapper from register_kernel (L346-362):

def backend_impl(*args, **kwargs):
    result = self._backend_fns[device_type](*args, **kwargs)    # user's fn
    def get_module():                                            # closure
        return inspect.getmodule(self._backend_fns[device_type])
    schema = self._opoverload._schema
    if not schema._is_view_op():
        utils._c_check_aliasing_constraint(                       # aliasing check
            self._name, args, kwargs, result, get_module)
    return result

Every call pays an aliasing-constraint check against self._opoverload._schema (iterates inputs/outputs, compares storage pointers) and a closure construction.

4. CustomOpDef.__call__ — L697. One extra Python frame when the op is invoked by the decorated name (e.g. nvfp4_gemm(x) in-module). Call sites that go through torch.ops.trtllm.nvfp4_gemm(x) bypass this frame.

What the low-level Library.define + impl path does

lib = Library("trtllm", "FRAGMENT")
lib.define(schema_str)       # declare op
lib.impl(name, fn, "CUDA")   # user's fn registered directly on the CUDA key

• No Autograd kernel registered → a call with requires_grad=True falls through to the dispatcher's C++-level "no autograd kernel" path (same user-facing error as before), no Python wrapper runs on the common inference path where no tensor needs grad.
• No ADInplaceOrView wrapper.
• CUDA key points at the user function directly; no backend_impl wrapper, no aliasing check on each call.

Call path comparison (same torch.ops.trtllm.my_op(x) invocation)

@custom_op path (pure, non-mutating op):

C++ dispatcher
 └─ Autograd key
     └─ autograd_impl(keyset, x)                        [Python]
         ├─ is_grad_enabled(); _any_requires_grad(*args)
         ├─ Metadata(keyset, kwargs)
         └─ forward_no_grad(x, Metadata)                [Python]
              └─ with _AutoDispatchBelowAutograd():     [Python ctx mgr]
                  └─ op.redispatch(..., x)              [re-dispatch]
                       └─ C++ dispatcher: CUDA key
                            └─ backend_impl(x)          [Python]
                                 ├─ user_fn(x)
                                 ├─ inspect.getmodule(...) closure
                                 └─ _c_check_aliasing_constraint(...)

Library.define + impl path:

C++ dispatcher
 └─ CUDA key
     └─ user_fn(x)

Two fewer Python frames + no aliasing check + no re-dispatch.

Microbenchmark

Isolated per-call cost on B200 + PyTorch 2.10, 20k-iter tight loop with x.clone() as a minimal kernel (tmp/bench_custom_op_overhead_v2.py):

*Pure dispatcher tax = total per-call minus the plain-Python-fn kernel floor.

Key observations:

• @fast_custom_op's hot path (torch.ops.trtllm.<name>(...)) is byte-identical to manual Library.define + impl at runtime: both resolve to the same OpOverload and skip all the Python wrappers. No wrapper overhead from the helper itself.
• The proxy __call__ path is marginally cheaper because the OpOverload object is cached at decoration time.

Estimated contribution on LTX-2 at baseline (@custom_op, 12us/call gross):

End-to-end measurements

Per-step host-time attribution, 1 GPU non-cuda-graph vs compile modes (from nsys, denoise_step NVTX range, baseline @custom_op state):

The E2E win scales with how much other host overhead (kernel launch APIs) has already been absorbed by torch.compile + CUDA Graph — on torch.compile + cuda_graph the @custom_op wrapper becomes a large fraction of what's left.

When to keep @custom_op

Keep @torch.library.custom_op for ops that (a) have non-empty mutates_args, (b) need autograd, or (c) are under active development and benefit from the richer Python-side error messages.

Functional equivalence

After the switch, torch.ops.trtllm.nvfp4_gemm and torch.ops.trtllm.tunable_fp4_quantize resolve to the same OpOverload objects with the same schema string as before — so eager Python, torch.compile, Dynamo, FX passes, and any downstream FX pattern matcher (e.g. ar_residual_norm that looks for torch.ops.trtllm.nvfp4_gemm.default) all see an identical op. The only difference is the dispatch path, which is shorter. register_fake still provides the same FakeTensor/meta shape inference used by torch.compile tracing.

Side-effect validation

Verified parity with @custom_op baseline via 13 targeted tests: numerical correctness, torch.compile (fullgraph + dynamic), FakeTensor/meta propagation, inference_mode / no_grad, mutation-safety (args not mutated, output not aliased to inputs), error cases (wrong dtype/device) raising identical RuntimeError, CPU-tensor dispatch failing identically, autograd-path with requires_grad=True raising the same error (both ops are non-differentiable by design).

E2E smoke test: 10-step LTX-2 1 GPU run with torch_compile=true using the @fast_custom_op form — exit code 0, steady-state 0.75s/step matches the manual-Library form.

Test Coverage

• tests/unittest/_torch/thop/test_nvfp4_gemm.py covers nvfp4_gemm — passes unchanged.
• tunable_fp4_quantize is exercised by all NVFP4 Linear / MoE tests and the LTX-2 integration.
• No new test surface: registration-API change only; behavior and schema are preserved.

PR Checklist

• PR description clearly explains what and why.

• PR Follows TRT-LLM CODING GUIDELINES.

• Test cases are provided for new code paths.

• Any new dependencies have been scanned for license and vulnerabilities.

• CODEOWNERS updated if ownership changes.

• Documentation updated as needed.

• Update tava architecture diagram if significant design change.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.